#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (C) 2023 reinovo, Inc. All Rights Reserved 
#
# @Time    : 2023/12/28 下午5:14
# @Author  : hmm
# @Email   : liuyuhang0531@foxmail.com
# @File    : Canera_Calibration.py
# @Software: PyCharm

import numpy as np
import cv2 as cv
import glob
import argparse

parser = argparse.ArgumentParser(description="相机标定程序")

parser.add_argument("--images_path", type=str, help="标定图像到位置")
parser.add_argument("--length", type=int, help="标定板的长边的角点个数")
parser.add_argument("--width", type=int, help="标定板的短边的角点个数")
parser.add_argument("--save_path", type=str,default="../data/camera_param.npz", help="标定板的短边的角点个数")

args = parser.parse_args()

length =  args.length
width = args.width
load_path = args.images_path
save_path = args.save_path

# 迭代终止条件
criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)
# 设置目标点：(0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((width*length,3), np.float32)
objp[:,:2] = np.mgrid[0:length,0:width].T.reshape(-1,2)
objp = objp
# 对象点和图像点列表
objpoints = [] # 真实世界空间的3d点
imgpoints = [] # 图像平面的2d点
images = glob.glob(path+'/*.jpg')
for fname in images:
    img = cv.imread(fname)
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    # 查找棋盘点
    ret, corners = cv.findChessboardCorners(gray, (length,width), None)
    # 如果找到棋盘点，加入对象点列表，图像列表
    if ret == True:
        objpoints.append(objp)
        corners2 = cv.cornerSubPix(gray,corners, (11,11), (-1,-1), criteria)
        imgpoints.append(corners2)
        # 绘制角点
        cv.drawChessboardCorners(img, (length,width), corners2, ret)
        cv.imshow('img', img)
        cv.waitKey(500)
cv.destroyAllWindows()


ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(objpoints,
                                                  imgpoints,
                                                  gray.shape[::-1],
                                                  None, None)

print("camera matrix :\n",mtx)
print("distortion coefficients :\n",dist)

np.savez("../data/camera_param.npz",mtx=mtx,dist=dist,rvecs=rvecs,tvecs=tvecs)

mean_error = 0
for i in range(len(objpoints)):
    imgpoints2,_ = cv.projectPoints(objpoints[i],rvecs[i],tvecs[i],mtx,dist)
    # 计算L2范数
    error = cv.norm(imgpoints[i],imgpoints2,cv.NORM_L2)/len(imgpoints2)
    mean_error += error

print("total error: {}".format(mean_error/len(objpoints)))